import numpy as np
from sklearn.linear_model import LogisticRegression
import pickle

# 读取傅里叶变换后的数据集
genres_list = ['classical', 'jazz', 'country', 'pop', 'rock', 'metal']

X = []
y = []

for g in genres_list:
    for n in range(100):
        rad = f"G:/akaa/genres/trainset/{g}.{str(n).zfill(5)}.fft.npy"
        fft_features = np.load(rad)
        X.append(fft_features)
        y.append(genres_list.index(g))

X = np.array(X)
y = np.array(y)

# 训练模型
model = LogisticRegression(solver='sag', multi_class='multinomial', max_iter=10000)
model.fit(X, y)

with open('model.pkl', 'wb') as f:
    pickle.dump(model, f)